import os
import cv2
import imagehash
from PIL import Image
import torch
from imagehash import ImageHash


class VideoFrameExtractor:
    def __init__(self, interval_sec=1.5, threshold=3, similarity_algo="ahash"):
        self.interval_sec = interval_sec
        self.threshold = threshold
        self.similarity_algo = similarity_algo
        self.image_extract_dict = {}

    def extract_similar_frames(self, video_path, output_frame_path):
        print(
            f"extract_similar_frames interval_sec {self.interval_sec} threshold {self.threshold}"
        )
        if not os.path.exists(output_frame_path):
            os.makedirs(output_frame_path)

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")

        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        interval = int(fps * self.interval_sec)

        try:
            index = 0
            for frame_count in range(0, total_frames, interval):
                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
                ret, frame = cap.read()
                if not ret:
                    break
                image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                self._extract_first_frame(image, frame_count, output_frame_path)
                # 进度更新
                progress = (frame_count / total_frames) * 100
                index = index + 1
                if index % 10 == 0:
                    print(f"处理进度: {progress:.2f}%")
        finally:
            cap.release()
            print("帧提取完成。")

    def _extract_first_frame(self, image, frame_num, output_frame_path):
        if self.similarity_algo == "ahash":
            # image_hash = self._average_hash(image)
            image_hash = imagehash.average_hash(image)
        elif self.similarity_algo == "phash":
            image_hash = imagehash.phash(image)
        elif self.similarity_algo == "dhash":
            image_hash = imagehash.dhash(image)
        elif self.similarity_algo == "whash":
            image_hash = imagehash.whash(image)
        else:
            raise ValueError("Invalid similarity algorithm")

        for _, image_info in self.image_extract_dict.items():
            if image_hash - image_info["hash"] <= self.threshold:
                return  # 发现相似图片，不再保存

        frame_filename = os.path.join(output_frame_path, f"{frame_num}.jpg")
        self.image_extract_dict[frame_filename] = {
            "hash": image_hash,
            "frame": frame_num,
        }
        image.save(frame_filename)
    
    def _average_hash(self, image, hash_size=8, mean=torch.mean):
        # type: (Image.Image, int, MeanFunc) -> ImageHash
        """
        Average Hash computation with GPU acceleration using PyTorch.

        @image must be a PIL instance.
        @mean how to determine the average luminescence. can try torch.median instead.
        """
        if hash_size < 2:
            raise ValueError('Hash size must be greater than or equal to 2')

        # reduce size and complexity, then convert to grayscale
        image = image.convert('L').resize((hash_size, hash_size), Image.LANCZOS)

        # find average pixel value; 'pixels' is an array of the pixel values, ranging from 0 (black) to 255 (white)
        pixels = torch.tensor(list(image.getdata()), dtype=torch.float32).view(hash_size, hash_size)
        
        # 如果有可用的GPU，则将数据移动到GPU上
        if torch.cuda.is_available():
            pixels = pixels.cuda()
        
        avg = mean(pixels)

        # create string of bits
        diff = pixels > avg
        # make a hash
        # 将结果移回CPU并转换为NumPy数组
        diff_numpy = diff.cpu().numpy()
        return ImageHash(diff_numpy)